import torch
import os
from torch.nn import init
import numpy as np

model_path = '/work_dir/pre_trained_weight/cascade_rcnn_s50_fpn_syncbn-backbone+head_mstrain-range_1x_coco_20201122_213640-763cc7b5.pth'
new_model_path = '/work_dir/pre_trained_weight'
num_classes = 8
channels = 4



def weight_modify():
    model_coco = torch.load(model_path)

    #4 channel
    model_coco['state_dict']['backbone.stem.0.weight'] = torch.cat( [model_coco['state_dict']['backbone.stem.0.weight']]  + [torch.unsqueeze(model_coco['state_dict']['backbone.stem.0.weight'][:, 0, :, :], 1)], dim=1)


    #8 classes
    model_coco["state_dict"]["roi_head.bbox_head.0.fc_cls.weight"] = model_coco["state_dict"][
                                                                         "roi_head.bbox_head.0.fc_cls.weight"][
                                                                     :num_classes, :]
    model_coco["state_dict"]["roi_head.bbox_head.1.fc_cls.weight"] = model_coco["state_dict"][
                                                                         "roi_head.bbox_head.1.fc_cls.weight"][
                                                                     :num_classes, :]
    model_coco["state_dict"]["roi_head.bbox_head.2.fc_cls.weight"] = model_coco["state_dict"][
                                                                         "roi_head.bbox_head.2.fc_cls.weight"][
                                                                     :num_classes, :]
    # bias
    model_coco["state_dict"]["roi_head.bbox_head.0.fc_cls.bias"] = model_coco["state_dict"][
                                                                       "roi_head.bbox_head.0.fc_cls.bias"][
                                                                   :num_classes]
    model_coco["state_dict"]["roi_head.bbox_head.1.fc_cls.bias"] = model_coco["state_dict"][
                                                                       "roi_head.bbox_head.1.fc_cls.bias"][
                                                                   :num_classes]
    model_coco["state_dict"]["roi_head.bbox_head.2.fc_cls.bias"] = model_coco["state_dict"][
                                                                       "roi_head.bbox_head.2.fc_cls.bias"][
                                                                   :num_classes]
    # save new model
    torch.save(model_coco, "/work_dir/pre_trained_weight/coco_pretrained_weights_classes_%d_inchannels_%d.pth" % (num_classes, channels))



if __name__ == '__main__':
    weight_modify()